import os
import torch
import numpy as np
from PIL import Image
import torch.nn.functional as F
from segment_anything import sam_model_registry, SamPredictor
from config import get_config
from logger import setup_logger

config = get_config()
logger = setup_logger('sam_model')

class SAMModel:
    """SAM模型封装类，用于加载模型和生成图像embedding"""
    def __init__(self):
        self.device = torch.device(config.DEVICE)
        self.model = None
        self.predictor = None
        self.model_loaded = False
        self.cache = {} if config.ENABLE_CACHE else None
    
    def load_model(self):
        """加载SAM模型"""
        if self.model_loaded:
            return True
        
        try:
            # 检查模型文件是否存在
            if not os.path.exists(config.SAM_MODEL_PATH):
                logger.error(f"模型文件不存在: {config.SAM_MODEL_PATH}")
                # 如果不存在，尝试从Segment Anything仓库下载
                logger.info("尝试从Segment Anything仓库下载模型...")
                # 创建models目录（如果不存在）
                os.makedirs(os.path.dirname(config.SAM_MODEL_PATH), exist_ok=True)
                
                # 根据模型类型确定下载链接
                model_download_urls = {
                    'vit_h': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth',
                    'vit_l': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth',
                    'vit_b': 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth',
                    'mobile_sam': 'https://github.com/ChaoningZhang/MobileSAM/releases/download/v1.0/mobile_sam.pth'
                }
                
                if config.MODEL_TYPE in model_download_urls:
                    download_url = model_download_urls[config.MODEL_TYPE]
                    logger.info(f"开始下载模型: {download_url}")
                    
                    try:
                        # 使用requests下载模型
                        import requests
                        from tqdm import tqdm
                        
                        response = requests.get(download_url, stream=True)
                        total_size = int(response.headers.get('content-length', 0))
                        block_size = 1024
                        
                        with open(config.SAM_MODEL_PATH, 'wb') as f, tqdm(
                            desc=os.path.basename(config.SAM_MODEL_PATH),
                            total=total_size, 
                            unit='iB', 
                            unit_scale=True,
                            unit_divisor=1024,
                        ) as bar:
                            for data in response.iter_content(block_size):
                                bar.update(len(data))
                                f.write(data)
                        
                        if total_size != 0 and bar.n != total_size:
                            logger.error("模型下载失败: 文件大小不匹配")
                            return False
                        
                        logger.info(f"模型下载成功: {config.SAM_MODEL_PATH}")
                    except Exception as e:
                        logger.error(f"下载模型失败: {e}")
                        return False
                else:
                    logger.error(f"不支持的模型类型: {config.MODEL_TYPE}")
                    return False
            
            
            logger.info(f"加载SAM模型: {config.SAM_MODEL_PATH}，设备: {self.device}")
            # 加载SAM模型
            self.model = sam_model_registry[config.MODEL_TYPE](checkpoint=config.SAM_MODEL_PATH)
            self.model.to(self.device)
            
            # 如果有CUDA，尝试使用Torch.compile进行优化
            if self.device.type == 'cuda' and hasattr(torch, 'compile'):
                try:
                    logger.info("使用Torch.compile优化模型...")
                    self.model = torch.compile(self.model)
                except Exception as e:
                    logger.warning(f"Torch.compile优化失败: {e}")
            
            # 创建预测器
            self.predictor = SamPredictor(self.model)
            self.model_loaded = True
            logger.info("SAM模型加载成功")
            return True
        except Exception as e:
            logger.error(f"加载SAM模型失败: {e}")
            return False
    
    def preprocess_image(self, image):
        """预处理图像"""
        if isinstance(image, str):
            # 从文件路径加载图像
            image = Image.open(image).convert('RGB')
        elif isinstance(image, bytes):
            # 从字节流加载图像
            from io import BytesIO
            image = Image.open(BytesIO(image)).convert('RGB')
        elif not isinstance(image, Image.Image):
            raise ValueError("输入图像必须是文件路径、字节流或PIL Image对象")
        
        # 调整图像大小
        if config.IMAGE_RESIZE > 0:
            width, height = image.size
            max_dim = max(width, height)
            scale = config.IMAGE_RESIZE / max_dim
            new_width = int(width * scale)
            new_height = int(height * scale)
            image = image.resize((new_width, new_height), Image.LANCZOS)
        
        return np.array(image)
    
    def generate_embedding(self, image, image_id=None):
        """生成图像的embedding"""
        if not self.model_loaded:
            if not self.load_model():
                raise RuntimeError("无法加载SAM模型")
        
        # 检查缓存
        if self.cache and image_id and image_id in self.cache:
            logger.debug(f"从缓存获取图像embedding: {image_id}")
            return self.cache[image_id]
        
        try:
            # 预处理图像
            processed_image = self.preprocess_image(image)
            
            # 使用predictor生成embedding
            with torch.no_grad():
                # 设置图像，这一步会生成图像embedding
                self.predictor.set_image(processed_image)
                # 获取生成的embedding
                image_embedding = self.predictor.get_image_embedding().cpu().numpy()
            
            # 缓存结果
            if self.cache and image_id:
                # 如果缓存已满，移除最早的项目
                if len(self.cache) >= config.CACHE_SIZE:
                    self.cache.pop(next(iter(self.cache)))
                self.cache[image_id] = image_embedding
            
            logger.debug(f"成功生成图像embedding，形状: {image_embedding.shape}")
            return image_embedding
        except Exception as e:
            logger.error(f"生成图像embedding失败: {e}")
            raise
    
    def clear_cache(self):
        """清除缓存"""
        if self.cache:
            self.cache.clear()
            logger.info("缓存已清除")
    
    def to(self, device):
        """移动模型到指定设备"""
        if self.model_loaded:
            self.device = torch.device(device)
            self.model.to(self.device)
            logger.info(f"模型已移动到设备: {self.device}")
        return self

# 创建全局模型实例
sam_model = SAMModel()